import json
import os
import sys

import backtrader as bt
import numpy as np
import pandas as pd

from run_backtrader2 import (
    CustomFutureCommission,
    CustomStockCommission,
    MultiFactorStrategy,
    SSPandasData,
)


class CustomStockCanShortCommission(bt.CommInfoBase):
    """
    股票佣金计算
    """

    params = (
        ("commission", 0.01),  # 0.01%佣金比例
        ("margin", 0.60),  # 12%的保证金比例
        ("stocklike", False),  # 不是股票回测
        ("automargin", 0.60),  # 由backtrader内部逻辑决定的必须设置一个正数
        ("commtype", bt.CommInfoBase.COMM_PERC),  # 是按照比例交佣金
    )


class LongShortStrategy(MultiFactorStrategy):
    """
    选top_num支股票做多，top_num支股票做空
    """

    params = (
        ("unit_k", 1),  # 成交单位量，
        ("is_stock", False),  # 是否是股票回测
        ("basedir", None),  # 保存回测结果的路径
        ("top_num", 30),  # 每天选取排名靠前的30只股票
        ("max_trade_cent", 0.05),  # 每天每只股票最多成交市场成交量5%
        ("trade_pos", 0.5),  # 最多用现有资金的50%买股票
    )

    def next(self):
        if not hasattr(self, "sdatas"):
            self.sdatas = {
                x._name: x
                for x in filter(lambda x: not x._name.startswith("rong"), self.datas)
            }
            self.rongdatas = {
                x._name: x
                for x in filter(lambda x: x._name.startswith("rong"), self.datas)
            }

        # self.broker.getcommissioninfo(self.data)
        top_num = self.p.top_num
        max_trade_cent = self.p.max_trade_cent
        trade_pos = self.p.trade_pos
        now = self.datas[0].datetime.datetime()

        target_value = {
            data: data.pred[0]
            for data in self.sdatas.values()
            if not np.isnan(data.pred[0])
        }
        # 获取当天每只股票的因子值

        rank_target = sorted(target_value.items(), key=lambda x: x[1], reverse=True)
        # 对股票因子值从高到低排名

        choose_target = [x[0] for x in rank_target[:top_num]]
        # choose_target = []
        choose_short_target = [
            self.rongdatas[f"rong_{x[0]._name}"] for x in rank_target[-top_num:]
        ]
        # choose_short_target = []
        # 选取排名靠前的因子

        print(
            f"日期：{now} 选择：{sorted([x._name for x in choose_target + choose_short_target])}"
        )
        self.every_day_choose.append(
            {
                "datetime": now,
                "choose_stocks": ",".join(
                    [x._name for x in choose_target + choose_short_target]
                ),
            }
        )
        all_worth = self.broker.getvalue()
        short_position_dict = {
            data._name: self.broker.getcommissioninfo(data).p.margin
            for data in choose_short_target
        }
        short_position_margin_sum = sum(short_position_dict.values())
        target_pos_mm = {
            data._name: all_worth
            * trade_pos
            * 1
            / (len(choose_target) + short_position_margin_sum)
            // (data.close[0] * self.p.unit_k)
            for data in choose_target
        }

        target_pos_mm_short = {
            data._name: -(
                all_worth
                * trade_pos
                * 1
                / (len(choose_target) + short_position_margin_sum)
                // (data.close[0] * self.p.unit_k)
            )
            for data in choose_short_target
        }
        # 根据当前资产价值计算目标仓位

        target_pos_mm = {**target_pos_mm, **target_pos_mm_short}

        cur_pos = {
            data._name: self.getposition(data).size // self.p.unit_k
            for data in self.datas
        }
        modify_pos = {
            data: target_pos_mm.get(data._name, 0) - cur_pos.get(data._name, 0)
            for data in self.datas
        }
        # 根据当前仓位和目标仓位计算买卖股票数量

        for data, mv in modify_pos.items():
            if mv < 0:
                mv = min(int(data.volume[0] * max_trade_cent), -mv)
                self.sell(data, size=self.p.unit_k * int(mv))
        for data, mv in modify_pos.items():
            if mv > 0:
                mv = min(int(data.volume[0] * max_trade_cent), mv)
                self.buy(data, size=self.p.unit_k * int(mv))

        self.all_cash.append(
            {
                "datetime": self.datas[0].datetime.datetime(),
                "cash": self.broker.getvalue(),
            }
        )


def main_job(
    result_file: str,
    code_market_list: list,
    start: str,
    end: str,
    cash: float,
    slip: float,
    is_stock: bool,
):
    """
    :param result_file: 保存的模型预测结果
    :param code_market_list: 选择回测的股票代码，为空list则回测所有
    :param start: 回测时间范围“20210101”
    :param end: 结束时间范围“20240701”
    :param cash: 回测使用的金额
    :param slip: 回测滑点，0.0001表示0.01%的滑点
    :param is_stock: 是否回测股票
    :return:
    """
    dir_base = os.path.join(os.path.dirname(result_file), f"equity_with_long_short2")
    os.makedirs(dir_base, exist_ok=True)
    df = pd.read_csv(result_file)
    if "volume" not in df.columns:
        df["volume"] = 1000
    if "openInterest" not in df.columns:
        df["openinterest"] = 0.0
    else:
        df.rename(columns={"openInterest": "openinterest"}, inplace=True)
    df.drop("date_time", inplace=True, axis=1)
    df["datetime"] = pd.to_datetime(
        df["time"] + 8 * 60 * 60 * 1000, unit="ms", utc=False
    )  # 时区转换
    df.drop("time", inplace=True, axis=1)
    df.reset_index(drop=True, inplace=True)
    if code_market_list:
        df = df[df["code_market"].isin(code_market_list)].reset_index()
    df = df[
        (df["datetime"] >= pd.to_datetime(start))
        & (df["datetime"] < pd.to_datetime(end))
    ]
    df.set_index("datetime", inplace=True)
    df.sort_index(inplace=True)

    # 创建一个Cerebro引擎
    cerebro = bt.Cerebro()

    # 添加策略
    cerebro.addstrategy(
        LongShortStrategy,
        unit_k=100 if is_stock else 10,
        is_stock=is_stock,
        basedir=dir_base,
    )
    dt_idx = df.index.drop_duplicates()
    stock_data = {}
    for xcode, df_group in df.groupby("code_market"):
        # 添加数据
        df_new = df_group.reindex(dt_idx).ffill().bfill()
        df_new.sort_index(inplace=True)
        stock_data[xcode] = df_new
    for xcode, df_group in stock_data.items():
        sdf = df_group.sort_index()
        data_feed = SSPandasData(dataname=sdf)
        cerebro.adddata(data_feed, name=xcode)
        cmi = CustomStockCommission() if is_stock else CustomFutureCommission()
        cerebro.broker.addcommissioninfo(cmi, name=xcode)
        data_feed = SSPandasData(dataname=sdf)
        cerebro.adddata(data_feed, name=f"rong_{xcode}")
        cerebro.broker.addcommissioninfo(
            CustomStockCanShortCommission(), name=f"rong_{xcode}"
        )

    cerebro.broker.setcash(cash)

    cerebro.broker.set_slippage_perc(perc=slip)
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe_ratio")
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name="annual_return")

    start_cash = cerebro.broker.getvalue()
    results = cerebro.run()
    end_cash = cerebro.broker.getvalue()
    sharpe_ratio = results[0].analyzers.sharpe_ratio.get_analysis()
    max_drawdown = results[0].analyzers.drawdown.get_analysis()
    annual_return = results[0].analyzers.annual_return.get_analysis()

    summary = dict(
        result_file=result_file,
        start=start,
        end=end,
        cash=cash,
        start_cash=start_cash,
        end_cash=end_cash,
        sharpe_ratio=sharpe_ratio,
        max_drawdown=max_drawdown,
        annual_return=annual_return,
        cmi=str(cmi.p),
        code_market_list=code_market_list,
    )
    print(summary)
    with open(os.path.join(dir_base, "backtest.json"), "w") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    print("Backtest result saved as backtest_result.png")


if __name__ == "__main__":
    datafile, start, end = sys.argv[1:]
    main_job(
        datafile,
        [],
        start,
        end,
        10000000,
        slip=0.0002,
        is_stock=True,
    )
